from concurrent.futures import ProcessPoolExecutor, as_completed

import numpy as np
import torch
from radarModel.model import SimVP, QPENet
from utils.readData import read_npy_from_path
from utils.save import save_image, save_rain_npy, turn_dbz_to_zr, save_history
import threading
from utils.Logger import get_logger
import torch.nn.functional as F
from utils.rainDraw import draw_cnn_png
from utils.nameUtils import generate_future_npy_filenames
import time
import yaml

logger = get_logger()  # 配置日志


# 雷达回波外推
def predict_dual(file_path: list, config: dict):
    '''

    :param config: 配置文件
    :param file_path:从java传递的数据库中最新的10条npy文件名称
    :param prefix:保存npy文件所在的目录
    :return:
    '''
    # 加载模型
    logger.info("加载雷达回波模型")
    Radarmodel = SimVP(([10, 16, 100, 100]), 64, 256, 4, 8).to("cpu")
    states = torch.load(config["radarModelPath"], map_location='cpu')
    Radarmodel.load_state_dict(states['state_dict'])
    Radarmodel.eval()
    input = read_npy_from_path(file_path, config["radar"]["save_path"]).to("cpu")
    # 输出结果并且进行保存
    with torch.no_grad():
        output_60min = Radarmodel(input)
        output_120_min = Radarmodel(output_60min)
        # [1,20,1,400,400]
        output = torch.cat([output_60min, output_120_min], dim=1)
        logger.info("雷达回波外推结束，保存雷达回波预测结果")
        # 保存可视化图片
        threading.Thread(target=save_image, args=(output, config, file_path[-1])).start()
        # 利用Z-R进行降雨估计
        threading.Thread(target=predict_rain_with_zr, args=(output, file_path[-1], config)).start()

        # 利用CNN进行降雨估计
        # predict_rain(inputTensor=output[0], lastTime=file_path[-1], config=config)


# 降雨估计
def predict_rain(inputTensor: torch.Tensor, lastTime: str, config: dict):
    '''
    对于输入400*400的网格  首先填充为424*424  然后滑动切片分为25*25大小的窗口 每个窗口输入模型得到一个值  最终得到400*400个值  最后将这些值还原为预测的图像
    本质上是考虑到周边窗口内部dbz的影响
    :param lastTime: 表示当前时间戳  202403010403
    :param inputTensor:(20,1,400,400) 分别表示(T,C,H,W)
    :param config:
    :return: numpy->(20,400,400)
    '''
    # 加载模型
    logger.info("加载降雨模型")
    rainModel = QPENet().to("cpu")
    states = torch.load(config["rainModelPath"], map_location='cpu')
    rainModel.load_state_dict(states["states"])

    rainModel.eval()
    T_length, C, height, width = inputTensor.shape
    batch_size = config["image"]["cnn"]["batch_size"]
    pre_z_all = [None] * T_length
    start = time.time()
    name_list = generate_future_npy_filenames(config['image']['cnn']['save_path'], lastTime, "png", 20)
    npy_list = generate_future_npy_filenames(config['radar']['cnn_dbz_path'], lastTime, "npy", 20)
    futures = {}
    # 线程池提交处理  同时保持前后顺序
    with ProcessPoolExecutor() as executor:
        for t in range(T_length):
            future = executor.submit(predict_rain_single, inputTensor[t, :, :, :], rainModel, batch_size, name_list[t], npy_list[t])
            futures[future] = t
        # 等待任务完成并收集结果
        for future in as_completed(futures):
            index = futures[future]
            result = future.result()
            pre_z_all[index] = result

    end = time.time()
    print(f"降雨模型结束，模型推理时间为 {end - start:.2f}s")
    # 最终的预测结果形状为 (T_length, 400, 400)，即 20 帧 400x400 的结果
    pre_z_all = np.array(pre_z_all)
    return pre_z_all


def predict_rain_single(frame, model, batch_size, filename_png, filename_npy):
    '''
    单帧降水模型的推演
    :param filename_png:
    :param filename_npy:
    :param batch_size:
    :param frame:(C,H,W)
    :param model:
    :return:
    '''
    # 使用 F.pad 进行填充，填充12个像素，保持滑动窗口有效
    dbz_padding = F.pad(frame, (12, 12, 12, 12), mode='constant', value=0)
    windows = dbz_padding.unfold(1, 25, 1).unfold(2, 25, 1)  # 生成滑动窗口
    windows = windows.contiguous().view(-1, 1, 25, 25)  # 变为 (160000, 1, 25, 25)
    num_windows = windows.size(0)
    pre_z = []

    # 分批进行推理
    with torch.no_grad():
        for i in range(0, num_windows, batch_size):
            batch_windows = windows[i:i + batch_size]  # 按批次选择窗口
            batch_predictions = model(batch_windows)  # 批量推理
            pre_z.append(batch_predictions.cpu().data.numpy())

    # 合并所有批次的预测结果，并将结果重塑为 (400*400)
    pre_z = np.concatenate(pre_z, axis=0)
    pre_z = pre_z.reshape(400, 400)
    # 保存CNN图片和对应npy文件
    threading.Thread(target=draw_cnn_png, args=(pre_z, filename_png)).start()
    threading.Thread(target=save_rain_npy, args=(pre_z, filename_npy, config)).start()

    return pre_z


def predict_rain_with_zr(inputTensor: torch.Tensor, lastTime: str, config: dict):
    '''
    通过ZR关系将预测的DBZ保存为降雨量结果
    :param inputTensor:
    :param lastTime:
    :param config:
    :return:
    '''
    logger.info("开始进行降雨估计")
    data = inputTensor[0, :, 0, ...].cpu().numpy()
    data = np.clip(data, 0, 1)
    pngList = generate_future_npy_filenames(config['image']['cnn']['save_path'], lastTime, "png", 20)
    npyList = generate_future_npy_filenames(config['radar']['cnn_dbz_path'], lastTime, "npy", 20)
    backupList = generate_future_npy_filenames(config['radar']['cnn_backup'], lastTime, "npy", 20)

    # 多线程提交保存降雨结果和可视化
    with ProcessPoolExecutor() as executor:
        for t in range(20):
            executor.submit(turn_dbz_to_zr, data[t], npyList[t], pngList[t], backupList[t])
    # 保存历史降水记录
    save_history(data, lastTime)


if __name__ == '__main__':
    path = ["202406042330"] * 10
    with open('/Users/luofan/PyCharm/ShijiaZhuang/config.yaml', 'r') as file:
        config = yaml.safe_load(file)
    predict_dual(path, config)
